"""
Session management for CodeMCP Gateway.

Manages user sessions for multi-step analysis workflows,
maintaining state across requests and providing session persistence.
"""

import asyncio
import json
import uuid
from typing import Dict, Any, Optional, List, Set
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from enum import Enum

from ..core.config import Config
from ..core.error_handler import handle_error, log_info, log_debug, log_warning, CodeMCPError


class SessionStatus(Enum):
    """Session status types."""
    ACTIVE = "active"
    IDLE = "idle"
    EXPIRED = "expired"
    TERMINATED = "terminated"


@dataclass
class SessionState:
    """Represents the state of a user session."""
    session_id: str
    user_id: Optional[str] = None
    status: SessionStatus = SessionStatus.ACTIVE
    created_at: datetime = field(default_factory=datetime.utcnow)
    last_activity: datetime = field(default_factory=datetime.utcnow)
    expires_at: Optional[datetime] = None
    
    # Session data
    current_project: Optional[str] = None
    current_files: List[str] = field(default_factory=list)
    analysis_history: List[Dict[str, Any]] = field(default_factory=list)
    call_trees: Dict[str, Any] = field(default_factory=dict)  # file_path -> call_tree
    session_variables: Dict[str, Any] = field(default_factory=dict)
    
    # Manual exploration state
    manual_mode_active: bool = False
    current_node_path: Optional[List[str]] = None
    exploration_stack: List[Dict[str, Any]] = field(default_factory=list)
    bookmarks: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    
    # Preferences
    preferences: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert session state to dictionary."""
        return {
            "session_id": self.session_id,
            "user_id": self.user_id,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "last_activity": self.last_activity.isoformat(),
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
            "current_project": self.current_project,
            "current_files": self.current_files,
            "analysis_history_count": len(self.analysis_history),
            "call_trees_count": len(self.call_trees),
            "manual_mode_active": self.manual_mode_active,
            "current_node_path": self.current_node_path,
            "bookmarks_count": len(self.bookmarks),
            "preferences": self.preferences
        }
    
    def is_expired(self) -> bool:
        """Check if session is expired."""
        if self.expires_at is None:
            return False
        return datetime.utcnow() > self.expires_at
    
    def update_activity(self):
        """Update last activity timestamp."""
        self.last_activity = datetime.utcnow()
        if self.status == SessionStatus.IDLE:
            self.status = SessionStatus.ACTIVE


class SessionManager:
    """Manages user sessions for the gateway."""
    
    def __init__(self, config: Config):
        self.config = config
        self.sessions: Dict[str, SessionState] = {}
        self.user_sessions: Dict[str, Set[str]] = {}  # user_id -> set of session_ids
        
        # Session configuration
        self.session_timeout_minutes = 60  # Default session timeout
        self.max_sessions_per_user = 5
        self.max_total_sessions = 1000
        self.cleanup_interval_minutes = 15
        
        # Background tasks
        self.cleanup_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        # Statistics
        self.stats = {
            "total_sessions_created": 0,
            "active_sessions": 0,
            "expired_sessions": 0,
            "terminated_sessions": 0,
            "cleanup_runs": 0
        }
    
    async def initialize(self):
        """Initialize the session manager."""
        # Load configuration
        self._load_session_config()
        
        # Start cleanup task
        self.is_running = True
        self.cleanup_task = asyncio.create_task(self._cleanup_loop())
        
        log_info("Session manager initialized successfully")
    
    async def shutdown(self):
        """Shutdown the session manager."""
        self.is_running = False
        
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
        
        # Optionally save session data for persistence
        await self._save_sessions()
        
        log_info("Session manager shutdown complete")
    
    def _load_session_config(self):
        """Load session configuration from config."""
        # Use custom settings if available
        custom = self.config.custom_settings
        
        self.session_timeout_minutes = custom.get("session_timeout_minutes", 60)
        self.max_sessions_per_user = custom.get("max_sessions_per_user", 5)
        self.max_total_sessions = custom.get("max_total_sessions", 1000)
        self.cleanup_interval_minutes = custom.get("session_cleanup_interval_minutes", 15)
        
        log_debug(f"Session config: timeout={self.session_timeout_minutes}min, "
                 f"max_per_user={self.max_sessions_per_user}, "
                 f"max_total={self.max_total_sessions}")
    
    def create_session(self, user_id: str = None) -> SessionState:
        """
        Create a new session.
        
        Args:
            user_id: Optional user identifier
            
        Returns:
            New session state
            
        Raises:
            CodeMCPError: If session creation fails
        """
        # Check total session limit
        if len(self.sessions) >= self.max_total_sessions:
            # Try to clean up expired sessions first
            self._cleanup_expired_sessions()
            
            if len(self.sessions) >= self.max_total_sessions:
                raise CodeMCPError("Maximum number of sessions reached")
        
        # Check per-user session limit
        if user_id:
            user_session_count = len(self.user_sessions.get(user_id, set()))
            if user_session_count >= self.max_sessions_per_user:
                raise CodeMCPError(f"Maximum sessions per user ({self.max_sessions_per_user}) reached")
        
        # Generate session ID
        session_id = str(uuid.uuid4())
        
        # Calculate expiration time
        expires_at = datetime.utcnow() + timedelta(minutes=self.session_timeout_minutes)
        
        # Create session
        session = SessionState(
            session_id=session_id,
            user_id=user_id,
            expires_at=expires_at
        )
        
        # Store session
        self.sessions[session_id] = session
        
        # Track user sessions
        if user_id:
            if user_id not in self.user_sessions:
                self.user_sessions[user_id] = set()
            self.user_sessions[user_id].add(session_id)
        
        # Update statistics
        self.stats["total_sessions_created"] += 1
        self.stats["active_sessions"] += 1
        
        log_info(f"Created session {session_id} for user {user_id}")
        return session
    
    def get_session(self, session_id: str) -> Optional[SessionState]:
        """
        Get session by ID.
        
        Args:
            session_id: Session identifier
            
        Returns:
            Session state or None if not found
        """
        session = self.sessions.get(session_id)
        
        if session is None:
            return None
        
        # Check if expired
        if session.is_expired():
            self._expire_session(session_id)
            return None
        
        # Update activity
        session.update_activity()
        return session
    
    def update_session(self, session_id: str, updates: Dict[str, Any]) -> bool:
        """
        Update session data.
        
        Args:
            session_id: Session identifier
            updates: Dictionary of updates to apply
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        # Apply updates
        for key, value in updates.items():
            if hasattr(session, key):
                setattr(session, key, value)
            else:
                # Store in session variables
                session.session_variables[key] = value
        
        session.update_activity()
        log_debug(f"Updated session {session_id} with {len(updates)} changes")
        return True
    
    def add_analysis_to_history(self, 
                               session_id: str, 
                               analysis_type: str,
                               file_path: str,
                               result: Dict[str, Any]) -> bool:
        """
        Add analysis result to session history.
        
        Args:
            session_id: Session identifier
            analysis_type: Type of analysis performed
            file_path: File that was analyzed
            result: Analysis result
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        # Create history entry
        history_entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "analysis_type": analysis_type,
            "file_path": file_path,
            "result_summary": {
                "success": result.get("success", False),
                "processing_time": result.get("processing_time"),
                "error": result.get("error", {}).get("message") if result.get("error") else None
            }
        }
        
        session.analysis_history.append(history_entry)
        
        # Keep only last 100 entries to prevent memory issues
        if len(session.analysis_history) > 100:
            session.analysis_history = session.analysis_history[-100:]
        
        session.update_activity()
        return True
    
    def store_call_tree(self, 
                       session_id: str, 
                       file_path: str, 
                       call_tree: Dict[str, Any]) -> bool:
        """
        Store call tree in session.
        
        Args:
            session_id: Session identifier
            file_path: File path the call tree belongs to
            call_tree: Call tree data
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        session.call_trees[file_path] = {
            "call_tree": call_tree,
            "timestamp": datetime.utcnow().isoformat()
        }
        
        session.update_activity()
        return True
    
    def get_call_tree(self, session_id: str, file_path: str) -> Optional[Dict[str, Any]]:
        """
        Get stored call tree for a file.
        
        Args:
            session_id: Session identifier
            file_path: File path
            
        Returns:
            Call tree data or None
        """
        session = self.get_session(session_id)
        if not session:
            return None
        
        return session.call_trees.get(file_path, {}).get("call_tree")
    
    def enter_manual_mode(self, session_id: str, entry_point: Dict[str, Any]) -> bool:
        """
        Enter manual exploration mode.
        
        Args:
            session_id: Session identifier
            entry_point: Starting point for manual exploration
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        session.manual_mode_active = True
        session.current_node_path = [entry_point.get("name", "root")]
        session.exploration_stack = [entry_point]
        
        session.update_activity()
        log_debug(f"Session {session_id} entered manual exploration mode")
        return True
    
    def exit_manual_mode(self, session_id: str) -> bool:
        """
        Exit manual exploration mode.
        
        Args:
            session_id: Session identifier
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        session.manual_mode_active = False
        session.current_node_path = None
        session.exploration_stack.clear()
        
        session.update_activity()
        log_debug(f"Session {session_id} exited manual exploration mode")
        return True
    
    def navigate_to_node(self, 
                        session_id: str, 
                        node_info: Dict[str, Any]) -> bool:
        """
        Navigate to a specific node in manual exploration.
        
        Args:
            session_id: Session identifier
            node_info: Information about the target node
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session or not session.manual_mode_active:
            return False
        
        # Add to exploration stack
        session.exploration_stack.append(node_info)
        
        # Update current path
        node_name = node_info.get("name", "unknown")
        if session.current_node_path:
            session.current_node_path.append(node_name)
        else:
            session.current_node_path = [node_name]
        
        session.update_activity()
        return True
    
    def add_bookmark(self, 
                    session_id: str, 
                    bookmark_name: str, 
                    node_info: Dict[str, Any]) -> bool:
        """
        Add a bookmark for manual exploration.
        
        Args:
            session_id: Session identifier
            bookmark_name: Name for the bookmark
            node_info: Node information to bookmark
            
        Returns:
            True if successful, False if session not found
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        session.bookmarks[bookmark_name] = {
            "node_info": node_info,
            "created_at": datetime.utcnow().isoformat(),
            "current_path": session.current_node_path.copy() if session.current_node_path else []
        }
        
        session.update_activity()
        return True
    
    def terminate_session(self, session_id: str, reason: str = "manual") -> bool:
        """
        Terminate a session.
        
        Args:
            session_id: Session identifier
            reason: Reason for termination
            
        Returns:
            True if successful, False if session not found
        """
        session = self.sessions.get(session_id)
        if not session:
            return False
        
        # Update status
        session.status = SessionStatus.TERMINATED
        
        # Remove from active tracking
        if session.user_id and session.user_id in self.user_sessions:
            self.user_sessions[session.user_id].discard(session_id)
            if not self.user_sessions[session.user_id]:
                del self.user_sessions[session.user_id]
        
        # Update statistics
        self.stats["active_sessions"] -= 1
        self.stats["terminated_sessions"] += 1
        
        log_info(f"Terminated session {session_id}, reason: {reason}")
        return True
    
    def _expire_session(self, session_id: str):
        """Mark session as expired."""
        session = self.sessions.get(session_id)
        if session and session.status != SessionStatus.EXPIRED:
            session.status = SessionStatus.EXPIRED
            
            # Update statistics
            if session.status == SessionStatus.ACTIVE:
                self.stats["active_sessions"] -= 1
            self.stats["expired_sessions"] += 1
    
    def _cleanup_expired_sessions(self):
        """Clean up expired sessions."""
        expired_sessions = []
        
        for session_id, session in self.sessions.items():
            if session.is_expired() or session.status in [SessionStatus.EXPIRED, SessionStatus.TERMINATED]:
                expired_sessions.append(session_id)
        
        # Remove expired sessions
        for session_id in expired_sessions:
            session = self.sessions.pop(session_id, None)
            if session and session.user_id:
                if session.user_id in self.user_sessions:
                    self.user_sessions[session.user_id].discard(session_id)
                    if not self.user_sessions[session.user_id]:
                        del self.user_sessions[session.user_id]
        
        if expired_sessions:
            log_info(f"Cleaned up {len(expired_sessions)} expired sessions")
        
        return len(expired_sessions)
    
    async def _cleanup_loop(self):
        """Background cleanup loop."""
        while self.is_running:
            try:
                self._cleanup_expired_sessions()
                self.stats["cleanup_runs"] += 1
                
                # Wait for next cleanup interval
                await asyncio.sleep(self.cleanup_interval_minutes * 60)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                log_warning(f"Error in session cleanup loop: {e}")
                await asyncio.sleep(60)  # Wait longer on error
    
    async def _save_sessions(self):
        """Save sessions for persistence (placeholder)."""
        # In a production system, this would save to a database or persistent storage
        # For now, we'll just log the session count
        active_count = sum(1 for s in self.sessions.values() 
                          if s.status == SessionStatus.ACTIVE)
        log_info(f"Session manager shutdown: {active_count} active sessions")
    
    def get_session_statistics(self) -> Dict[str, Any]:
        """Get session statistics."""
        # Count sessions by status
        status_counts = {status.value: 0 for status in SessionStatus}
        for session in self.sessions.values():
            status_counts[session.status.value] += 1
        
        return {
            **self.stats,
            "current_sessions": len(self.sessions),
            "status_distribution": status_counts,
            "users_with_sessions": len(self.user_sessions),
            "session_timeout_minutes": self.session_timeout_minutes,
            "max_sessions_per_user": self.max_sessions_per_user,
            "max_total_sessions": self.max_total_sessions
        }
    
    def get_user_sessions(self, user_id: str) -> List[Dict[str, Any]]:
        """Get all sessions for a user."""
        if user_id not in self.user_sessions:
            return []
        
        user_session_ids = self.user_sessions[user_id]
        return [
            self.sessions[session_id].to_dict() 
            for session_id in user_session_ids 
            if session_id in self.sessions
        ]